# -*- coding: utf-8 -*-
"""
ibvs_analysis.py — 仅保存两张图、不显示窗口的健壮版
1) <csv_stem>_timeseries.png  —— 误差/速度/速度模 三子图（含 t_start / t_conv 标注）
2) 若提供 --mocap_file → <csv_stem>_trajectory.png（3D 轨迹）
   否则 → <csv_stem>_extra.png（res/avg 曲线）

用法：
  python ibvs_analysis.py data.csv --save_figures --no_show
  python ibvs_analysis.py data.csv --mocap_file mocap.csv --save_figures --no_show
  # 如需跳过前5行表头：
  python ibvs_analysis.py data.csv --skip_rows 5 --save_figures --no_show
"""

from __future__ import annotations
import argparse, math
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# ---------- small utils ----------
def _find_columns(df: pd.DataFrame, patterns: List[str]) -> List[str]:
    cols = []
    for pat in patterns:
        for c in df.columns:
            if pat.lower() in c.lower():
                cols.append(c)
    # 去重保持顺序
    seen = set(); out = []
    for c in cols:
        if c not in seen:
            seen.add(c); out.append(c)
    return out


def _detect_motion_start(
    v_norm: np.ndarray, times: np.ndarray,
    fraction: float = 0.05, consecutive: int = 10
) -> Tuple[Optional[int], Optional[float]]:
    if len(v_norm) < consecutive:
        return None, None
    max_val = np.nanmax(np.abs(v_norm))
    if not np.isfinite(max_val) or max_val <= 0:
        return None, None
    thresh = fraction * max_val
    for i in range(len(v_norm) - consecutive + 1):
        if np.all(v_norm[i:i+consecutive] > thresh):
            return i, float(times[i])
    return None, None


def _detect_convergence(
    err_norm: np.ndarray, times: np.ndarray, start_idx: int,
    threshold: float, settle_time: float
) -> Tuple[Optional[int], Optional[float]]:
    if settle_time <= 0 or len(times) < 2:
        return None, None
    dt = np.median(np.diff(times))
    if not np.isfinite(dt) or dt <= 0:
        return None, None
    need = int(max(1, math.ceil(settle_time / dt)))
    for i in range(start_idx, len(err_norm) - need):
        if np.all(err_norm[i:i+need] < threshold):
            return i, float(times[i])
    return None, None


def _read_csv_smart(path: Path, skip_rows: int = 0, encoding: str = "utf-8-sig") -> pd.DataFrame:
    """
    更健壮的 CSV 读取：
    - 先用 sep=None + engine='python' 自动嗅探
    - 失败则分别尝试 [',', '\\t', ';', '\\s+']
    - 支持跳过表头/注释行：--skip_rows
    """
    kwargs = dict(encoding=encoding, skiprows=skip_rows)
    try:
        return pd.read_csv(path, sep=None, engine="python", **kwargs)  # 自动嗅探
    except Exception:
        for sep in [",", "\t", ";", r"\s+"]:
            try:
                return pd.read_csv(path, sep=sep, engine="python", **kwargs)
            except Exception:
                continue
    # 最后一次尝试：给出更友好提示
    raise ValueError(
        f"无法解析数据文件：{path}\n"
        "建议：检查分隔符（逗号/制表符/分号），或使用 --skip_rows 跳过表头。"
    )


# ---------- main analysis ----------
def analyze_run(
    csv_file: str,
    mocap_file: Optional[str] = None,
    start_threshold: float = 0.05,
    start_consecutive: int = 10,
    error_threshold: float = 5.0,
    settle_time: float = 0.5,
    plot: bool = True,
    save_figures: bool = False,
    output_dir: Optional[str] = None,
    show: bool = True,
    skip_rows: int = 0,
    encoding: str = "utf-8-sig",
) -> Dict[str, float]:
    csv_path = Path(csv_file)
    if not csv_path.is_file():
        raise FileNotFoundError(f"数据文件未找到: {csv_file}")

    # 友好检查：防止把 .py 当数据传入
    if csv_path.suffix.lower() in {".py", ".ipynb"}:
        raise ValueError(
            f"检测到传入的是代码文件 {csv_path.name}。"
            " 请把“CSV 数据文件”作为第一个参数传入，而不是 .py。"
        )

    # 读取数据（更健壮）
    df = _read_csv_smart(csv_path, skip_rows=skip_rows, encoding=encoding)

    # 时间列
    time_cols = _find_columns(df, ["t_ros", "time", "timestamp", "t_sec"])
    if not time_cols:
        raise ValueError("CSV 中未找到时间列（尝试包含 't_ros'/'time'/'timestamp'/'t_sec' 的列名）")
    time_col = time_cols[0]
    times = df[time_col].to_numpy(dtype=float)

    # 速度指令列
    vx_cols = _find_columns(df, ["vx", "twist_x", "cmd_vx", "lx", "twist.lx"])
    vy_cols = _find_columns(df, ["vy", "twist_y", "cmd_vy", "ly", "twist.ly"])
    vz_cols = _find_columns(df, ["vz", "twist_z", "cmd_vz", "lz", "twist.lz"])
    if not (vx_cols and vy_cols and vz_cols):
        raise ValueError("未找到速度指令列（期望 vx/vy/vz 或 twist.* / cmd_v* / l* 等列）")

    vx = df[vx_cols[0]].to_numpy(dtype=float)
    vy = df[vy_cols[0]].to_numpy(dtype=float)
    vz = df[vz_cols[0]].to_numpy(dtype=float)
    v_norm = np.linalg.norm(np.vstack([vx, vy, vz]), axis=0)

    # 误差列（优先 res*，再 pid_e*，再 err*；排除 avg）
    res_cols = [c for c in df.columns if ("res" in c.lower()) and ("avg" not in c.lower())]
    pid_cols = [c for c in df.columns if ("pid_e" in c.lower()) and ("avg" not in c.lower())]
    generic_cols = [c for c in df.columns if ("err" in c.lower()) and ("avg" not in c.lower())]
    err_cols = res_cols or pid_cols or generic_cols
    if not err_cols:
        err_cols = _find_columns(df, ["ex", "ey", "ez"])
    if not err_cols:
        raise ValueError("未找到误差列（如 res*/pid_e* 或 ex/ey/ez）")

    err_matrix = df[err_cols].to_numpy(dtype=float)
    err_norm = np.linalg.norm(err_matrix, axis=1)

    # 运动开始 & 收敛
    start_idx, t_start = _detect_motion_start(v_norm, times, fraction=start_threshold, consecutive=start_consecutive)
    if t_start is None:
        t_start = float("nan")
    if start_idx is not None:
        conv_idx, t_conv = _detect_convergence(err_norm, times, start_idx, error_threshold, settle_time)
        if t_conv is None:
            t_conv = float("nan")
    else:
        conv_idx, t_conv = None, float("nan")

    T_conv = (t_conv - t_start) if (np.isfinite(t_start) and np.isfinite(t_conv)) else float("nan")
    metrics = {"t_start": t_start, "t_conv": t_conv, "T_conv": T_conv}

    if plot:
        times_rel = times - times[0]

        # 图1：误差/速度/速度模
        fig_ts, axes_ts = plt.subplots(3, 1, sharex=True, figsize=(8, 9))
        for col in err_cols:
            axes_ts[0].plot(times_rel, df[col].to_numpy(dtype=float), label=col)
        axes_ts[0].set_ylabel("Error (pixels)")
        axes_ts[0].set_title("Pixel Error vs Time")
        axes_ts[0].legend(loc="best")

        axes_ts[1].plot(times_rel, vx, label=vx_cols[0])
        axes_ts[1].plot(times_rel, vy, label=vy_cols[0])
        axes_ts[1].plot(times_rel, vz, label=vz_cols[0])
        axes_ts[1].set_ylabel("Commanded velocity (m/s)")
        axes_ts[1].set_title("Commanded Velocity vs Time")
        axes_ts[1].legend(loc="best")

        axes_ts[2].plot(times_rel, v_norm, label="||v_cmd||")
        axes_ts[2].set_ylabel("Velocity norm (m/s)")
        axes_ts[2].set_xlabel("Time (s)")
        axes_ts[2].set_title("Velocity Norm vs Time")
        axes_ts[2].legend(loc="best")

        for ax in axes_ts:
            if np.isfinite(t_start):
                ax.axvline(t_start - times[0], linestyle="--", linewidth=1)
                ax.text(t_start - times[0], ax.get_ylim()[1], " t_start", rotation=90, va="top", ha="right")
            if np.isfinite(t_conv):
                ax.axvline(t_conv - times[0], linestyle="-.", linewidth=1)
                ax.text(t_conv - times[0], ax.get_ylim()[1], " t_conv", rotation=90, va="top", ha="right")
        fig_ts.tight_layout()

        # 图2候选A：res/avg
        fig_extra = None
        res_cols_all = _find_columns(df, ["res"])
        avg_cols = _find_columns(df, ["avg"])
        if res_cols_all or avg_cols:
            fig_extra, axes_extra = plt.subplots(2, 1, sharex=True, figsize=(8, 6))
            if res_cols_all:
                for col in res_cols_all:
                    axes_extra[0].plot(times_rel, df[col].to_numpy(dtype=float), label=col)
                axes_extra[0].set_ylabel("Residual")
                axes_extra[0].set_title("Residual Signals vs Time")
                axes_extra[0].legend(loc="best")
            else:
                axes_extra[0].set_visible(False)

            if avg_cols:
                for col in avg_cols:
                    axes_extra[1].plot(times_rel, df[col].to_numpy(dtype=float), label=col)
                axes_extra[1].set_ylabel("Average")
                axes_extra[1].set_title("Averaged Signals vs Time")
                axes_extra[1].legend(loc="best")
            else:
                axes_extra[1].set_visible(False)

            for ax in axes_extra:
                if np.isfinite(t_start):
                    ax.axvline(t_start - times[0], linestyle="--", linewidth=1)
                if np.isfinite(t_conv):
                    ax.axvline(t_conv - times[0], linestyle="-.", linewidth=1)
            fig_extra.tight_layout()

        # 图2候选B：3D 轨迹
        fig_3d = None
        if mocap_file:
            mpath = Path(mocap_file)
            if mpath.is_file():
                mdf = _read_csv_smart(mpath, skip_rows=0, encoding=encoding)
                ee_cols = _find_columns(mdf, ["ee_x", "ee_y", "ee_z", "arm_x", "arm_y", "arm_z"])
                face_cols = _find_columns(mdf, ["face_x", "face_y", "face_z", "target_x", "target_y", "target_z"])
                if len(ee_cols) >= 3 and len(face_cols) >= 3:
                    ee_x, ee_y, ee_z = (mdf[ee_cols[0]].to_numpy(float),
                                        mdf[ee_cols[1]].to_numpy(float),
                                        mdf[ee_cols[2]].to_numpy(float))
                    face_x, face_y, face_z = (mdf[face_cols[0]].to_numpy(float),
                                              mdf[face_cols[1]].to_numpy(float),
                                              mdf[face_cols[2]].to_numpy(float))
                    fig_3d = plt.figure(figsize=(7, 6))
                    ax3d = fig_3d.add_subplot(111, projection="3d")  # type: ignore
                    ax3d.plot(ee_x, ee_y, ee_z, label="End-effector path")
                    ax3d.plot(face_x, face_y, face_z, label="Target path")
                    ax3d.scatter([ee_x[0], ee_x[-1]], [ee_y[0], ee_y[-1]], [ee_z[0], ee_z[-1]], s=20)
                    ax3d.set_xlabel("X (m)"); ax3d.set_ylabel("Y (m)"); ax3d.set_zlabel("Z (m)")
                    ax3d.set_title("3D Trajectory of End-Effector and Target")
                    ax3d.legend(loc="best")
                    fig_3d.tight_layout()

        # ——仅保存两张图——
        if save_figures:
            outdir = Path(output_dir) if output_dir else csv_path.parent
            outdir.mkdir(parents=True, exist_ok=True)
            base = csv_path.stem

            p1 = outdir / f"{base}_timeseries.png"
            fig_ts.savefig(p1)

            if fig_3d is not None:
                p2 = outdir / f"{base}_trajectory.png"
                fig_3d.savefig(p2)
            elif fig_extra is not None:
                p2 = outdir / f"{base}_extra.png"
                fig_extra.savefig(p2)

            print("Saved figures:")
            print("  ", p1)
            if (fig_3d is not None) or (fig_extra is not None):
                print("  ", p2)

        if show:
            plt.show()
        else:
            plt.close('all')

    return metrics


def main() -> None:
    parser = argparse.ArgumentParser(description="分析 IBVS 实验 CSV 文件并仅保存两张图（不显示窗口）。")
    parser.add_argument("csv_file", help="CSV 数据文件路径（不是 .py）")
    parser.add_argument("--mocap_file", default=None, help="可选：动捕 CSV（若提供则第二张图为3D轨迹）")
    parser.add_argument("--start_threshold", type=float, default=0.05, help="运动开始阈值（速度模相对最大值比例）")
    parser.add_argument("--start_consecutive", type=int, default=10, help="运动开始需连续超过阈值的样本数")
    parser.add_argument("--error_threshold", type=float, default=5.0, help="收敛判据：误差模阈值")
    parser.add_argument("--settle_time", type=float, default=0.5, help="收敛保持时间（秒）")
    parser.add_argument("--no_plot", action="store_true", help="不生成图，仅计算数值")
    parser.add_argument("--save_figures", action="store_true", help="保存图像到磁盘")
    parser.add_argument("--output_dir", help="保存目录（默认与 CSV 同目录）")
    parser.add_argument("--no_show", action="store_true", help="仅保存、不显示窗口（切换 Agg 后端）")
    parser.add_argument("--skip_rows", type=int, default=0, help="跳过表头/注释行数，例如 5")
    parser.add_argument("--encoding", default="utf-8-sig", help="文件编码，默认 utf-8-sig（可改为 gbk 等）")
    args = parser.parse_args()

    # 仅保存不显示 → 切 Agg 后端
    if args.no_show:
        try:
            plt.switch_backend("Agg")
        except Exception:
            pass

    # 友好报错：误把 .py 当数据
    p = Path(args.csv_file)
    if p.suffix.lower() in {".py", ".ipynb"}:
        raise SystemExit(
            f"检测到传入的是代码文件：{p.name}\n"
            "请把 CSV 数据文件作为第一个参数，例如：\n"
            "  python ibvs_analysis.py .\\ibvs_D_10.10_T_15_28.csv --save_figures --no_show"
        )

    analyze_run(
        csv_file=args.csv_file,
        mocap_file=args.mocap_file,
        start_threshold=args.start_threshold,
        start_consecutive=args.start_consecutive,
        error_threshold=args.error_threshold,
        settle_time=args.settle_time,
        plot=not args.no_plot,
        save_figures=args.save_figures,
        output_dir=args.output_dir,
        show=(not args.no_show),
        skip_rows=args.skip_rows,
        encoding=args.encoding,
    )


if __name__ == "__main__":
    main()
